"""
### 名词
- 偏导数（Partial Derivatives）:偏导数用于描述多元函数中，函数对其中一个自变量的变化率（其他自变量保持不变）。
- 梯度（Gradients）:梯度是偏导数的向量形式，表示函数在所有自变量方向上的变化率。
- 自动微分（Automatic Differentiation，简称 AutoDiff）:自动微分是一种计算导数的算法，它通过追踪计算过程并生成
- 反向传播（Backpropagation）的链式法则，自动计算函数的梯度。
- 偏置单元（Bias Unit）

在大语言模型（LLM）中，自动微分（Automatic Differentiation，简称 AutoDiff） 是训练过程的核心技术之一，
它能够高效计算模型参数对损失函数的梯度，为基于梯度下降的参数优化提供基础。与传统手动推导导数的方式不同，自动
微分通过算法自动追踪计算过程并生成导数表达式，极大简化了复杂模型（如包含数十亿参数的 LLM）的训练实现。

### 核心原理：计算图与链式法则
自动微分的本质是将模型的计算过程拆解为计算图（Computation Graph），再通过链式法则反向传播梯度。

1. 计算图的构建
LLM 的前向计算（如输入文本→词嵌入→Transformer 层→输出概率分布）可以拆分为一系列基本运算（矩阵乘法、激活函数、
注意力计算等）。自动微分框架（如 PyTorch、TensorFlow）会在模型前向传播时动态构建计算图，其中：
- 节点表示张量（输入数据、模型参数、中间结果）
- 边表示运算（如+、×、softmax、LayerNorm等）

2. 反向传播求梯度
训练时，先通过前向传播计算损失函数（如交叉熵损失），再从损失函数出发，沿计算图反向遍历，根据链式法则逐层计算每个
参数的梯度（损失对参数的偏导数）。例如，对于计算链 L <- y <- z <- w (L为损失，w为参数），梯度计算为：
dL/dw = dL/dy * dy/dz * dz/dw

### 在 LLM 中的关键作用

LLM 的核心是通过海量数据调整参数（如注意力权重、线性层权重等），使模型输出更接近真实标签。自动微分在此过程中承担以
下关键任务：

1. 高效计算海量参数的梯度
LLM 通常包含数十亿甚至数万亿参数（如 GPT-3 有 1750 亿参数），手动推导所有参数的梯度表达式是不可能的。自动微分通过
计算图动态追踪参数与损失的依赖关系，仅对参与计算的参数求导，大幅降低了工程复杂度。

2. 支持复杂结构的梯度计算
LLM 中的 Transformer 模块包含多头注意力、残差连接、层归一化等复杂结构，自动微分能够无缝处理这些嵌套运算的梯度传递，
无需人工干预。例如，注意力权重的梯度会通过 softmax 的导数、矩阵乘法的导数等逐层反向传播。

3. 动态控制计算精度与效率
自动微分框架支持混合精度计算（如 FP16/FP8），在 LLM 训练中可显著降低内存占用和计算量；同时支持梯度裁剪（Gradient
Clipping），缓解训练中因梯度爆炸导致的模型不稳定问题。

### 与 LLM 训练流程的结合

以典型的 LLM 训练步骤为例，自动微分的作用体现在：
1. 前向传播：输入文本通过词嵌入层、Transformer 块等生成预测分布，框架同步构建计算图，记录参数与中间结果的依赖关系。
2. 损失计算：通过交叉熵损失函数（如CrossEntropyLoss）计算预测分布与真实标签的差异。
3. 反向传播：调用自动微分接口（如 PyTorch 的loss.backward()），框架沿计算图反向计算所有可训练参数的梯度。
4. 参数更新：优化器（如 Adam）使用自动微分得到的梯度，按规则更新参数（如 \(w = w - \eta \cdot \nabla L\)，\(\eta\) 为学习率）。

### LLM 中自动微分的挑战与优化

1.内存消耗：LLM 的计算图包含海量中间结果（如每个 Transformer 层的激活值），反向传播时需存储这些值以计算梯度。实际
中通过激活检查点（Activation Checkpointing） 优化：只保存部分关键激活值，其他值在反向传播时重新计算，以时间换内存。

2. 分布式训练适配：在数据并行或模型并行中，自动微分需支持梯度的跨设备聚合（如allreduce操作），框架（如 PyTorch DDP）已内置这些功能。

3. 数值稳定性：LLM 的深度结构可能导致梯度消失或爆炸，自动微分框架通过梯度缩放（Gradient Scaling）等技术缓解这一问题。

### 总结

自动微分是 LLM 训练的 “幕后引擎”，它通过计算图和链式法则自动完成复杂的梯度计算，使研究者能够专注于模型结构设计而
非导数推导。在 PyTorch 等框架中，这一过程被高度封装，开发者只需调用backward()即可触发梯度计算，极大推动了 LLM
的快速迭代与落地。
"""
import torch
import torch.nn.functional as F
from torch.autograd import grad

y = torch.tensor([1.0])
x1 = torch.tensor([1.1])
w1 = torch.tensor([2.2], requires_grad=True)
b = torch.tensor([0.0], requires_grad=True)

z = x1 * w1 + b
a = torch.sigmoid(z)

loss = F.binary_cross_entropy(a, y)

# manually calculate the gradient of loss with respect to w1 and b
# default, pytorch destroy the computation graph, so we need to set retain_graph=True to reuse the graph shortly
grad_L_w1 = grad(loss, w1, retain_graph=True)
grad_L_b = grad(loss, b, retain_graph=True)
print("grad_L_w1: ", grad_L_w1)
print("grad_L_b: ", grad_L_b)

# backward() will auto compute the gradient of all leaf nodes in the graph
loss.backward()
print("w1 grad: ", w1.grad)
print("b grad: ", b.grad)